import os

import numpy as np
import tensorflow as tf


class AutoRec(object):

    def __init__(self, visibleDimensions,
                 epochs=200, hiddenDimensions=50,
                 learningRate=0.1,
                 batchSize=100,
                 model_path='autorec_model'):
        self.epochs = epochs
        self.hiddenDimensions = hiddenDimensions
        self.learningRate = learningRate
        self.batchSize = batchSize
        self.visibleDimensions = visibleDimensions

        # 添加维度标识到模型路径, 防止数据维度不同导致新旧模型冲突
        self.model_path = f"{model_path}_dim{visibleDimensions}"

        # Load the model if it exists, otherwise build a new one
        if os.path.exists(os.path.join(model_path, 'saved_model.pb')):
            self.model = tf.keras.models.load_model(model_path)
            print("Model loaded from disk.")
        else:
            self.model = self.build_model()
            print("New model created.")

    def save_model(self):
        if not os.path.exists(self.model_path):
            os.makedirs(self.model_path)
        self.model.save(self.model_path)
        print(f"Model saved to {self.model_path}")

    def build_model(self):
        # Input layer
        input_layer = tf.keras.layers.Input(shape=(self.visibleDimensions,))

        # Hidden layer
        hidden = tf.keras.layers.Dense(self.hiddenDimensions, activation='sigmoid')(input_layer)

        # Output layer
        output_layer = tf.keras.layers.Dense(self.visibleDimensions, activation='sigmoid')(hidden)

        # Create the model
        model = tf.keras.models.Model(inputs=input_layer, outputs=output_layer)

        # Compile the model
        model.compile(optimizer=tf.keras.optimizers.RMSprop(learning_rate=self.learningRate),
                      loss='mean_squared_error')

        return model

    def Train(self, X):
        npX = np.array(X)

        # Train the model
        self.model.fit(npX, npX, epochs=self.epochs, batch_size=self.batchSize, verbose=1)

        # Save the model after training
        self.save_model()

    def GetRecommendations(self, inputUser):
        # Feed through a single user and return predictions from the output layer.
        inputUser = np.array(inputUser)
        rec = self.model.predict(inputUser)

        return rec[0]

# Example usage:
# visibleDimensions = 1000  # Example number of items
# autoRec = AutoRec(visibleDimensions)
# autoRec.Train(X)  # X is your training data
# recommendations = autoRec.GetRecommendations(inputUser)  # inputUser is a single user's data
